import numpy as np

def average_growth_coefficient(X, U, V, alpha=0.03):
    """
    平均增长系数法迭代实现交通分布预测

    >>> average_growth_coefficient(X, U, V, alpha=0.03)

    输入参数
    ----------
    X : numpy.ndarray
        当前年各个交通小区的 OD 矩阵
    U : numpy.ndarray
        未来年的 O_i，行向量
    V : numpy.ndarray
        未来年的 D_j，行向量
    alpha : float, optional
        收敛准则, by default 0.03

    返回值
    -------
    numpy.ndarray
        迭代后的OD表，仅包含OD部分
    """

    OD_mat = X

    future_Oi = U
    future_Dj = V

    iter_OD_mat = X

    iter_Oi = np.sum(OD_mat, axis=1) # 横向求和
    iter_Dj = np.sum(OD_mat, axis=0) # 纵向求和

    F_o = future_Oi / iter_Oi
    F_d = future_Dj / iter_Dj

    while not ( # 正式开始迭代
        (
            np.all(
                np.abs(F_o - 1) < alpha)
         ) and (
            np.all(
                np.abs(F_d - 1) < alpha)
                )):
        # 迭代停止的标准是未来年出行量与本年出行量，
        # 以及未来年吸引量与本年吸引量之间的比例小于
        # 收敛准则 alpha，即两者的比值减去 1 小于 
        # alpha
    
        f_mat = np.zeros(X.shape)
        
        for i in range(X.shape[0]):
            for j in range(X.shape[1]):
                f_mat[i][j] = 0.5 * (F_o[i] + F_d[j])
    
        iter_OD_mat = f_mat * iter_OD_mat
    
        iter_Oi = np.sum(iter_OD_mat, axis=1)
        iter_Dj = np.sum(iter_OD_mat, axis=0)
        
        # 计算新的 F_o 和 F_d

        F_o = future_Oi / iter_Oi
        F_d = future_Dj / iter_Dj
    
    return iter_OD_mat

def detroit(X, U, V, alpha=0.03):
    """
    底特律法迭代实现交通分布预测

    >>> detroit(X, U, V, alpha=0.03)
    
    考虑将来年的出行分布不仅与出行的发生和吸引增长率有关,还与出行
    生成量的增长率有关,考虑的因素较平均增长系数方法全面,但同样是
    收敛速度慢,需要多次迭代才能求得将来年的分布交通量。

    输入参数
    ----------
    X : numpy.ndarray
        当前年各个交通小区的 OD 矩阵
    U : numpy.ndarray
        未来年的 O_i，行向量
    V : numpy.ndarray
        未来年的 D_j，行向量
    alpha : float, optional
        收敛准则, by default 0.03

    返回值
    -------
    numpy.ndarray
        迭代后的OD表，仅包含OD部分
    """

    OD_mat = X

    future_Oi = U
    future_Dj = V

    future_sum = np.sum(future_Dj)

    iter_OD_mat = X

    iter_Oi = np.sum(OD_mat, axis=1)
    iter_Dj = np.sum(OD_mat, axis=0)

    iter_sum = np.sum(iter_Oi)

    F_o = future_Oi / iter_Oi
    F_d = future_Dj / iter_Dj

    while not ( # 正式开始迭代
        (
            np.all(
                np.abs(F_o - 1) < alpha)
         ) and (
            np.all(
                np.abs(F_d - 1) < alpha)
                )):

        # rg_rate 是倒增长率 reciprocal growth rate
        # T^m / X
        rg_rate = (iter_sum / future_sum)
    
        f_mat = np.outer(F_o, F_d) * rg_rate
    
        iter_OD_mat = f_mat * iter_OD_mat
    
        iter_Oi = np.sum(iter_OD_mat, axis=1)
        iter_Dj = np.sum(iter_OD_mat, axis=0)

        iter_sum = np.sum(iter_Oi) 
    
    return iter_OD_mat

def frator(X, U, V, alpha=0.03):
    """
    福莱特法迭代实现交通分布预测
    
    佛莱特(Frator)方法较平均增长系数法收敛速度快，在满足相同的精
    度条件下迭代次数也少，因此在实际工作中广泛应用。


    >>> detroit(X, U, V, alpha=0.03)

    输入参数
    ----------
    X: numpy.ndarray
        当前年各个交通小区的 OD 矩阵
    U : numpy.ndarray
        未来年的 O_i，行向量
    V : numpy.ndarray
        未来年的 D_j，行向量
    alpha : float, optional
        收敛准则, by default 0.03

    返回值
    -------
    numpy.ndarray
        迭代后的OD表，仅包含OD部分
    """

    OD_mat = X

    future_Oi = U
    future_Dj = V

    iter_OD_mat = X

    iter_Oi = np.sum(OD_mat, axis=1)
    iter_Dj = np.sum(OD_mat, axis=0)

    F_o = future_Oi / iter_Oi
    F_d = future_Dj / iter_Dj

    while not ( # 正式开始迭代
        (
            np.all(
                np.abs(F_o - 1) < alpha)
         ) and (
            np.all(
                np.abs(F_d - 1) < alpha)
                )):

        Li = iter_Oi / (F_d @ X.T)
        Lj = iter_Dj / (X.T @ F_o)

        f_mat = np.zeros(X.shape)

        for i in range(iter_OD_mat.shape[0]):
            for j in range(iter_OD_mat.shape[1]):
                f_mat[i][j] = (
                    0.5 * (
                        Li[i] + Lj[j]
                        ) * F_o[i] * F_d[j])
    
        iter_OD_mat = f_mat * iter_OD_mat
    
        iter_Oi = np.sum(iter_OD_mat, axis=1)
        iter_Dj = np.sum(iter_OD_mat, axis=0)
        
        # 计算新的 F_o 和 F_d

        F_o = future_Oi / iter_Oi
        F_d = future_Dj / iter_Dj
    
    return iter_OD_mat